import asyncio
import aiohttp
import logging
from datetime import datetime
import socket
import time

class SSEDiagnostic:
    def __init__(self):
        self.logger = logging.getLogger("SSE_Diagnostic")
        self.logger.setLevel(logging.INFO)
        formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
        
        # 控制台输出
        ch = logging.StreamHandler()
        ch.setFormatter(formatter)
        self.logger.addHandler(ch)
        
        # 文件输出
        fh = logging.FileHandler('sse_diagnostic.log')
        fh.setFormatter(formatter)
        self.logger.addHandler(fh)
    
    async def test_sse_connection(self, url):
        """测试SSE连接并收集诊断信息"""
        self.logger.info("=" * 60)
        self.logger.info(f"开始SSE诊断测试 - 时间: {datetime.now().isoformat()}")
        self.logger.info(f"连接URL: {url}")
        
        results = {
            "start_time": datetime.now().isoformat(),
            "url": url,
            "events_received": 0,
            "event_types": set(),
            "errors": 0,
            "connection_duration": 0,
            "raw_data": []
        }
        
        try:
            # 检查端口是否开放
            host, port = self._parse_url(url)
            if not self._check_port(host, port):
                self.logger.error(f"端口 {port} 未开放")
                results["errors"] += 1
                return results
            
            # 测试HTTP端点
            http_url = f"http://{host}:{port}/api/health"
            http_success = await self._test_http_endpoint(http_url)
            if not http_success:
                self.logger.warning("HTTP端点测试失败")
                results["errors"] += 1
            
            # 测试SSE连接
            start_time = time.time()
            async with aiohttp.ClientSession() as session:
                async with session.get(url) as response:
                    # 检查响应头
                    content_type = response.headers.get('Content-Type', '')
                    if 'text/event-stream' not in content_type:
                        self.logger.warning(f"内容类型不是text/event-stream: {content_type}")
                    
                    # 读取事件流
                    async for line in response.content:
                        if len(results["raw_data"]) < 10:  # 只保存前10行
                            results["raw_data"].append(line.decode().strip())
                        
                        # 简单的事件计数
                        if line.startswith(b'data:'):
                            results["events_received"] += 1
                        elif line.startswith(b'event:'):
                            event_type = line.decode().split(':')[1].strip()
                            results["event_types"].add(event_type)
            
            results["connection_duration"] = time.time() - start_time
            self.logger.info(f"连接持续时间: {results['connection_duration']:.2f}秒")
            self.logger.info(f"收到事件数量: {results['events_received']}")
            self.logger.info(f"事件类型: {', '.join(results['event_types']) if results['event_types'] else '无'}")
            
        except asyncio.TimeoutError:
            self.logger.error("连接超时")
            results["errors"] += 1
        except aiohttp.ClientConnectorError as e:
            self.logger.error(f"连接错误: {str(e)}")
            results["errors"] += 1
        except Exception as e:
            self.logger.error(f"未知错误: {str(e)}", exc_info=True)
            results["errors"] += 1
        finally:
            results["end_time"] = datetime.now().isoformat()
            return results
    
    def _parse_url(self, url):
        """解析URL获取主机和端口"""
        from urllib.parse import urlparse
        parsed = urlparse(url)
        host = parsed.hostname or "localhost"
        port = parsed.port or 80
        return host, port
    
    def _check_port(self, host, port):
        """检查端口是否开放"""
        try:
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                s.settimeout(2)
                result = s.connect_ex((host, port))
                return result == 0
        except Exception:
            return False
    
    async def _test_http_endpoint(self, url):
        """测试HTTP端点是否可用"""
        try:
            async with aiohttp.ClientSession() as session:
                async with session.get(url, timeout=5) as response:
                    if response.status == 200:
                        self.logger.info("HTTP端点测试成功")
                        return True
                    self.logger.warning(f"HTTP端点返回状态码: {response.status}")
                    return False
        except Exception as e:
            self.logger.warning(f"HTTP端点测试失败: {str(e)}")
            return False
    
    def generate_report(self, results):
        """生成诊断报告"""
        report = [
            "=" * 60,
            "SSE连接诊断报告",
            "=" * 60,
            f"测试开始时间: {results['start_time']}",
            f"测试结束时间: {results['end_time']}",
            f"连接URL: {results['url']}",
            f"连接持续时间: {results.get('connection_duration', 0):.2f}秒",
            f"收到事件数量: {results['events_received']}",
            f"事件类型: {', '.join(results['event_types']) if results['event_types'] else '无'}",
            f"错误数量: {results['errors']}",
            ""
        ]
        
        if results["errors"] > 0:
            report.append("诊断发现问题:")
            if results.get("connection_duration", 0) == 0:
                report.append("- 无法建立连接")
            elif results["events_received"] == 0:
                report.append("- 未收到任何事件")
            report.append("")
        
        if results["raw_data"]:
            report.append("原始响应数据示例:")
            report.append("-" * 60)
            for i, line in enumerate(results["raw_data"]):
                report.append(f"{i+1}: {line}")
            report.append("-" * 60)
        
        report.append("建议解决方案:")
        if results["errors"] == 0:
            report.append("- 连接正常，请检查前端实现")
        else:
            report.append("- 检查后端服务是否运行")
            report.append("- 验证SSE端点实现是否正确")
            report.append("- 检查防火墙设置")
            report.append("- 测试不同浏览器")
        
        report.append("=" * 60)
        return "\n".join(report)

async def main():
    diagnostic = SSEDiagnostic()
    
    # 测试URL - 替换为您的SSE端点
    test_url = "http://localhost:8000/api/chat_stream?user_input=你好"
    
    results = await diagnostic.test_sse_connection(test_url)
    report = diagnostic.generate_report(results)
    
    print(report)
    diagnostic.logger.info("诊断完成")

if __name__ == "__main__":
    asyncio.run(main())